const faceapi = require("@vladmandic/face-api");
const path = require("path");
const fs = require("fs");

window.addEventListener("DOMContentLoaded", async () => {
  const video = document.getElementById("video");
  const labelsDir = path.join(__dirname, "labels");

  // 1. 加载所有模型
  console.log("正在加载模型...");
  await Promise.all([
    faceapi.nets.tinyFaceDetector.loadFromUri(path.join(__dirname, "models")),
    faceapi.nets.faceLandmark68Net.loadFromUri(path.join(__dirname, "models")),
    faceapi.nets.faceRecognitionNet.loadFromUri(path.join(__dirname, "models")),
    faceapi.nets.ssdMobilenetv1.loadFromUri(path.join(__dirname, "models")), // 识别需要更精确的模型
  ]);
  console.log("模型加载完毕");

  // 2. 加载带标签的参考图片并创建 FaceMatcher
  console.log("正在学习已知人脸...");
  const faceMatcher = new faceapi.FaceMatcher(await loadLabeledImages(), 0.6);
  console.log("人脸学习完毕，准备就绪");

  // 3. 启动摄像头
  navigator.mediaDevices
    .getUserMedia({ video: {} })
    .then((stream) => {
      video.srcObject = stream;
    })
    .catch((err) => console.error("获取摄像头失败:", err));

  video.addEventListener("play", () => {
    const canvas = faceapi.createCanvasFromMedia(video);
    document.querySelector(".container").append(canvas);
    const displaySize = { width: video.width, height: video.height };
    faceapi.matchDimensions(canvas, displaySize);

    setInterval(async () => {
      const detections = await faceapi
        .detectAllFaces(video)
        .withFaceLandmarks()
        .withFaceDescriptors();
      const resizedDetections = faceapi.resizeResults(detections, displaySize);

      // 清除上一帧的绘制内容
      canvas.getContext("2d").clearRect(0, 0, canvas.width, canvas.height);

      const results = resizedDetections.map((d) =>
        faceMatcher.findBestMatch(d.descriptor)
      );

      results.forEach((result, i) => {
        const box = resizedDetections[i].detection.box;
        // 使用 face-api.js 的 DrawBox 来绘制带标签的框
        const drawBox = new faceapi.draw.DrawBox(box, {
          label: result.toString(),
        });
        drawBox.draw(canvas);
      });
    }, 100);
  });

  /**
   * 加载 labels 文件夹中的图片，为每张图片计算人脸描述符
   */
  async function loadLabeledImages() {
    const labels = fs
      .readdirSync(labelsDir)
      .filter((file) => fs.statSync(path.join(labelsDir, file)).isDirectory());

    return Promise.all(
      labels.map(async (label) => {
        const descriptions = [];
        const labelPath = path.join(labelsDir, label);
        const imageFiles = fs.readdirSync(labelPath);

        for (const file of imageFiles) {
          const imgPath = path.join(labelPath, file);
          // faceapi.fetchImage 需要一个 URL，对于本地文件，需要转换一下格式
          const img = await faceapi.fetchImage(`file://${imgPath}`);
          const detections = await faceapi
            .detectSingleFace(img)
            .withFaceLandmarks()
            .withFaceDescriptor();
          if (detections) {
            descriptions.push(detections.descriptor);
          }
        }
        return new faceapi.LabeledFaceDescriptors(label, descriptions);
      })
    );
  }
});
